﻿#include "utils.h"
#include "openssl/rsa.h"
#include "openssl/pem.h"

#include <winsock2.h>
#include <QDebug>
#include <QPushButton>
#include <sstream>
#include <ctime>
#include <regex>

void utils_rsa::GenerateRsaKeyCstr(char* &prikey_cstr, char* &pubkey_cstr, int key_len, size_t& pubkey_len, size_t& prikey_len){
    prikey_len = 0, pubkey_len = 0;
    RSA *keypair = RSA_generate_key(key_len, RSA_3, NULL, NULL);
    BIO *pri_bio  = BIO_new(BIO_s_mem());
    BIO *pub_bio = BIO_new(BIO_s_mem());

    PEM_write_bio_RSAPrivateKey(pri_bio, keypair, NULL, NULL, 0, NULL, NULL);
    PEM_write_bio_RSAPublicKey(pub_bio, keypair);
    prikey_len = BIO_pending(pri_bio);
    pubkey_len = BIO_pending(pub_bio);
    prikey_cstr = new char[prikey_len+1];
    pubkey_cstr = new char[pubkey_len+1];
    BIO_read(pri_bio, prikey_cstr, prikey_len);
    BIO_read(pub_bio, pubkey_cstr, pubkey_len);
    prikey_cstr[prikey_len] = 0x00;
    pubkey_cstr[pubkey_len] = 0x00;

    RSA_free(keypair);
    BIO_free_all(pri_bio);
    BIO_free_all(pub_bio);
}


void utils_rsa::GenerateRsaKey(std::string &prikey, std::string &pubkey, int key_len){
    char* prikey_cstr = nullptr;
    char* pubkey_cstr = nullptr;
    size_t prikey_len = 0;
    size_t pubkey_len = 0;

    GenerateRsaKeyCstr(prikey_cstr, pubkey_cstr, key_len, pubkey_len, prikey_len);

    prikey = std::string(prikey_cstr, prikey_len);
    pubkey = std::string(pubkey_cstr, pubkey_len);

    delete[] prikey_cstr;
    delete[] pubkey_cstr;
}


bool utils_rsa::RsaEncryptUsingPubkey(const std::string pubkey, const std::string clear_text, std::string &encrypted_text){
    encrypted_text = "";
    int key_len = 0, block_len = 0, ret = 0, pos = 0;

    BIO* key_bio = BIO_new_mem_buf((unsigned char*)pubkey.c_str(), -1);
    RSA* rsa = RSA_new();

    rsa = PEM_read_bio_RSAPublicKey(key_bio, &rsa, NULL, NULL);
    if(!rsa){
        return false;
    }

    key_len = RSA_size(rsa);
    block_len = key_len - 11;

    char* sub_c_str = new char[key_len+1];
    memset(sub_c_str, 0x00, key_len+1);
    std::string sub_str;

    while(pos < clear_text.length()){
        sub_str = clear_text.substr(pos, block_len);
        memset(sub_c_str, 0x00, key_len+1);
        ret = RSA_public_encrypt(sub_str.size(), (const unsigned char*)sub_str.c_str(), (unsigned char*)sub_c_str, rsa, RSA_PKCS1_PADDING);
        if(ret >= 0){
            encrypted_text.append(std::string(sub_c_str, ret));
        }else{
            return false;
        }
        pos += block_len;
    }

    RSA_free(rsa);
    BIO_free(key_bio);
    delete[] sub_c_str;
    return true;
}


bool utils_rsa::RsaDecryptUsingPrikey(const std::string prikey, const std::string encrypted_text, std::string &clear_text){
    clear_text = "";
    std::string decrypted_text = "";
    int enc_text_len = 0, ret = 0, pos = 0;
    std::string sub_str = "";
    int block_len = 0;

    BIO* key_bio = BIO_new_mem_buf((unsigned char*)prikey.c_str(), -1);
    RSA* rsa = RSA_new();
    rsa = PEM_read_bio_RSAPrivateKey(key_bio, &rsa, NULL, NULL);
    if(rsa==nullptr){
        return false;
    }

    enc_text_len = RSA_size(rsa);
    block_len = enc_text_len - 11;

    char* target_c_str = new char[block_len+1];

    while(pos < sizeof(encrypted_text.c_str())){
        sub_str = encrypted_text.substr(pos, enc_text_len);
        memset(target_c_str, 0x00, block_len+1);
        ret = RSA_private_decrypt(enc_text_len, (const unsigned char*)sub_str.c_str(), (unsigned char*)target_c_str, rsa, RSA_PKCS1_PADDING);
        if(ret >= 0){
            decrypted_text.append(std::string(target_c_str, ret));
        }else{
            RSA_free(rsa);
            BIO_free(key_bio);
            delete[] target_c_str;
            return false;
        }
        pos += enc_text_len;
    }

    RSA_free(rsa);
    BIO_free(key_bio);
    delete[] target_c_str;
    clear_text = decrypted_text;
    return true;
}


bool utils_rsa::RsaEncryptUsingPrikey(const std::string prikey, const std::string clear_text, std::string &encrypted_text){
    encrypted_text = "";
    int key_len = 0, block_len = 0, ret = 0, pos = 0;

    BIO* key_bio = BIO_new_mem_buf((unsigned char*)prikey.c_str(), -1);
    RSA* rsa = RSA_new();

    rsa = PEM_read_bio_RSAPrivateKey(key_bio, &rsa, NULL, NULL);
    if(!rsa){
        return false;
    }

    key_len = RSA_size(rsa);
    block_len = key_len - 11;

    char* sub_c_str = new char[key_len+1];
    memset(sub_c_str, 0x00, key_len+1);
    std::string sub_str;

    while(pos < clear_text.length()){
        sub_str = clear_text.substr(pos, block_len);
        memset(sub_c_str, 0x00, key_len+1);
        ret = RSA_private_encrypt(sub_str.length(), (const unsigned char*)sub_str.c_str(), (unsigned char*)sub_c_str, rsa, RSA_PKCS1_PADDING);
        if(ret >= 0){
            encrypted_text.append(std::string(sub_c_str, ret));
        }else{
            return false;
        }
        pos += block_len;
    }

    RSA_free(rsa);
    BIO_free(key_bio);
    delete[] sub_c_str;
    return true;
}


bool utils_rsa::RsaDecryptUsingPubkey(const std::string pubkey, const std::string encrypted_text, std::string &clear_text){
    clear_text = "";
    std::string decrypted_text = "";
    int enc_text_len = 0, ret = 0, pos = 0;
    std::string sub_str = "";
    int block_len = 0;

    BIO* key_bio = BIO_new_mem_buf((unsigned char*)pubkey.c_str(), -1);
    RSA* rsa = RSA_new();
    rsa = PEM_read_bio_RSAPublicKey(key_bio, &rsa, NULL, NULL);
    if(rsa==nullptr){
        return false;
    }

    enc_text_len = RSA_size(rsa);
    block_len = enc_text_len - 11;

    char* target_c_str = new char[block_len+1];

    while(pos < sizeof(encrypted_text.c_str())){
        sub_str = encrypted_text.substr(pos, enc_text_len);
        memset(target_c_str, 0x00, block_len+1);
        ret = RSA_public_decrypt(enc_text_len, (const unsigned char*)sub_str.c_str(), (unsigned char*)target_c_str, rsa, RSA_PKCS1_PADDING);
        if(ret >= 0){
            decrypted_text.append(std::string(target_c_str, ret));
        }else{
            RSA_free(rsa);
            BIO_free(key_bio);
            delete[] target_c_str;
            return false;
        }
        pos += enc_text_len;
    }

    RSA_free(rsa);
    BIO_free(key_bio);
    delete[] target_c_str;
    clear_text = decrypted_text;
    return true;
}

std::string utils_rsa::StrToHexStr(std::string str){
    std::string ret;
    char* str_per_byte = new char[8];
    for(int i = 0; i < str.length(); i++){
        sprintf(str_per_byte, "%02x", (unsigned char)(str.c_str())[i]);
        ret.append(str_per_byte);
    }
    ret.append("\0");
    delete[] str_per_byte;
    return ret;
}


utils_sock::sock_fd_t utils_sock::CreateClientSockFd(std::string ip, unsigned short port){
    unsigned int sock_fd = 0;

    struct sockaddr_in addr{};
    size_t len = sizeof(addr);

    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = inet_addr(ip.c_str());

    sock_fd = socket(AF_INET, SOCK_STREAM, 0);
    int ret = connect(sock_fd, reinterpret_cast<struct sockaddr*>(&addr), len);
    if(ret == SOCKET_ERROR){
        closesocket(sock_fd);
        sock_fd = 0;
    }

    return sock_fd;
}


bool utils_sock::SendMsg(sock_fd_t sock_fd, std::string msg){
    int len = msg.size();
    int ret = send(sock_fd, msg.c_str(), len, 0);
    if(ret == SOCKET_ERROR){
        return false;
    } else {
        return true;
    }
}


bool utils_sock::SendRsaPubkeyEncryptedMsg(sock_fd_t sock_fd, std::string pubkey, std::string clear_msg){
    std::string encrypted_msg;
    if(!utils_rsa::RsaEncryptUsingPubkey(pubkey, clear_msg, encrypted_msg)){
        return false;
    } else {
        return utils_sock::SendMsg(sock_fd, encrypted_msg);
    }
}

utils_rsa::KeyType utils_rsa::CheckKeyType(std::string probable_key){
    BIO *bio = nullptr;
    RSA* rsa = nullptr;

    bio = BIO_new_mem_buf(probable_key.c_str(), probable_key.size());
    rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    if( rsa != nullptr ){
        RSA_free(rsa);
        BIO_free_all(bio);
        return RSA_PUBKEY_TYPE;
    }else{
        BIO_reset(bio);
        rsa = PEM_read_bio_RSAPrivateKey(bio, NULL, NULL, NULL);
        if( rsa != nullptr ){
            RSA_free(rsa);
            BIO_free_all(bio);
            return RSA_PRIKEY_TYPE;
        }else{
            RSA_free(rsa);
            BIO_free_all(bio);
            return OTHER;
        }
    }
}

void utils_log::QdebugWithLog(std::string new_content, std::string &log_str){
    qDebug()<<new_content.c_str();
    log_str.append(new_content);
}

void utils_log::UserOpLog(int op_code, std::string &log_str, ...){
    va_list args;
    va_start(args, log_str);
    if(op_code == utils_log::CLICK_PUSH_BUTTON){
        QPushButton* button;
        button = va_arg(args, QPushButton*);
        std::string detail = va_arg(args, const char*);
        std::stringstream ss;
        time_t tm = std::time(nullptr);
        ss << "User clicked button at time [" << tm << "]\n";
        ss << "Button name: [" << button->accessibleName().toUtf8().constData() << "]\n";
        ss << "Button address: [0x" << std::hex << button << "]\n";
        ss << "Detail: " << detail << "\n";
        qDebug()<<ss.str().c_str();
        button = nullptr;
    }
    va_end(args);
}

void utils_log::SystemInfoLog(std::string info_str, std::string &log_str){

}

void utils_log::SystemErrLog(int err_code, std::string &log_str, ...){
    va_list args;
}

bool utils_sock::CheckIpAddrFormat(std::string ip_addr){
    std::regex rg1("((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)");
    if(std::regex_match(ip_addr, rg1)){
        return true;
    }
    return false;
}

utils_sock::sock_fd_t utils_sock::CreateListenSockFd(unsigned short port){
    sock_fd_t sock_fd = 0;

    struct sockaddr_in addr{};
    size_t len = sizeof(addr);

    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = INADDR_ANY;

    sock_fd = socket(AF_INET, SOCK_STREAM, 0);
    int ret = connect(sock_fd, reinterpret_cast<struct sockaddr*>(&addr), len);
    if(bind(sock_fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr) ) == -1 ){

    }

    return sock_fd;
}

utils_sock::TcpSocketWrapper::TcpSocketWrapper(int type, unsigned int local_port, const char* remote_ip, unsigned int remote_port){
    this->type = type;
    this->local_port = local_port;
    this->remote_ip = remote_ip;
    this->remote_port = remote_port;
    this->recv_buffer = new char[512];
    this->fd = 0;
}

void utils_sock::TcpSocketWrapper::Connect(){
    if(type!=SOCK_TYPE_CLIENT){
        return;
    }
#if WIN32
    SetupWSA();
#endif
    this->fd = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(remote_port);
    addr.sin_addr.s_addr = inet_addr(remote_ip.c_str());
}

void utils_sock::TcpSocketWrapper::Listen(){

}

int utils_sock::TcpSocketWrapper::CheckStatus(){
    return status;
}
void utils_sock::TcpSocketWrapper::Disconnect(){}

void utils_sock::TcpSocketWrapper::Send(const char* buffer){}

utils_sock::TcpSocketWrapper::~TcpSocketWrapper(){}
#if WIN32
void utils_sock::TcpSocketWrapper::SetupWSA(){
    WORD wsa_ver = MAKEWORD(2, 2);
    struct WSAData wsa_data;
    WSAStartup(wsa_ver, &wsa_data);
}
void utils_sock::TcpSocketWrapper::ShutdownWSA(){
    WSACleanup();
}
#endif
